package com.sheng.leetcode.year2023.month03.day16;

import org.junit.Test;

/**
 * @author liusheng
 * @date 2023/03/16
 * <p>
 * 2488. 统计中位数为 K 的子数组<p>
 * <p>
 * 给你一个长度为 n 的数组 nums ，该数组由从 1 到 n 的 不同 整数组成。另给你一个正整数 k 。<p>
 * 统计并返回 nums 中的 中位数 等于 k 的非空子数组的数目。<p>
 * 注意：<p>
 * 数组的中位数是按 递增 顺序排列后位于 中间 的那个元素，如果数组长度为偶数，则中位数是位于中间靠 左 的那个元素。<p>
 * 例如，[2,3,1,4] 的中位数是 2 ，[8,4,3,5,1] 的中位数是 4 。<p>
 * 子数组是数组中的一个连续部分。<p>
 * <p>
 * 示例 1：<p>
 * 输入：nums = [3,2,1,4,5], k = 4<p>
 * 输出：3<p>
 * 解释：中位数等于 4 的子数组有：[4]、[4,5] 和 [1,4,5] 。<p>
 * <p>
 * 示例 2：<p>
 * 输入：nums = [2,3,1], k = 3<p>
 * 输出：1<p>
 * 解释：[3] 是唯一一个中位数等于 3 的子数组。<p>
 * <p>
 * 提示：<p>
 * n == nums.length<p>
 * 1 <= n <= 10^5<p>
 * 1 <= nums[i], k <= n<p>
 * nums 中的整数互不相同<p>
 * <p>
 * 来源：力扣（LeetCode）<p>
 * 链接：<a href="https://leetcode.cn/problems/count-subarrays-with-median-k">2488. 统计中位数为 K 的子数组</a><p>
 * 著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。<p>
 */
public class LeetCode2488 {

    @Test
    public void test01() {
        int[] nums = {3, 2, 1, 4, 5};
        int k = 4;
//        int[] nums = {2, 3, 1};
//        int k = 3;
        System.out.println(new Solution().countSubarrays(nums, k));
    }
}

class Solution {
    public int countSubarrays(int[] nums, int k) {
        int n = nums.length;
        int i = 0;
        // 获取数组 nums 中 k 元素所在的位置 i
        for (; nums[i] != k; ++i) {
        }
        // 创建一个长度为 2 * n + 1 的数组 cnt
        int[] cnt = new int[n << 1 | 1];
        // 返回值，初始值为 k 单独一个子数组
        int ans = 1;
        // 从 i 位置开始，不断地向右遍历，x 表示数组遍历过程中比 k 大的元素的数量
        int x = 0;
        for (int j = i + 1; j < n; ++j) {
            // 判断当前元素是否大于中位数 k
            x += nums[j] > k ? 1 : -1;
            // 如果满足条件，代表从 i 到 j 的子数组是满足题目要求的子字符串
            if (x >= 0 && x <= 1) {
                // 返回值自增
                ++ans;
            }
            ++cnt[x + n];
        }
        // 从 i 位置开始往左遍历
        x = 0;
        for (int j = i - 1; j >= 0; --j) {
            x += nums[j] > k ? 1 : -1;
            // 判断是否满足条件
            if (x >= 0 && x <= 1) {
                // 满足则自增
                ++ans;
            }
            // 是否存在跨越左右数组的子数组
            ans += cnt[-x + n] + cnt[-x + 1 + n];
        }
        return ans;
        // 暴力破解超时
//        int count = 0;
//        for (int i = 0; i < n; i++) {
//            for (int j = i; j < n; j++) {
//                int[] ints = Arrays.copyOfRange(nums, i, j + 1);
//                Arrays.sort(ints);
//                if (ints.length % 2 == 0) {
//                    if (ints[ints.length / 2 - 1] == k) {
//                        count++;
//                    }
//                } else {
//                    if (ints[ints.length / 2] == k) {
//                        count++;
//                    }
//                }
//            }
//        }
//        return count;
    }
}
